import { ChatOllama } from "@langchain/ollama";
import {
  ChatPromptTemplate,
  SystemMessagePromptTemplate,
} from "@langchain/core/prompts";
import { StringOutputParser } from "@langchain/core/output_parsers";
import {
  RunnableBranch,
  RunnableLambda,
  RunnablePassthrough,
  RunnableSequence,
} from "@langchain/core/runnables";

// 1. 模型
const model = new ChatOllama({
  model: "llama3",
});

// 2. 解析器
const parser = new StringOutputParser();

// 3. 三条“领域链” —— 物理 / 数学 / 其他
const physicsChain = ChatPromptTemplate.fromTemplate(
  `
你是一位物理学家，擅长使用中文回答物理相关的问题，当你不知道问题的答案时，你就回答不知道。
具体问题如下：
{input}
`
)
  .pipe(model)
  .pipe(parser); // 物理链

const mathChain = ChatPromptTemplate.fromTemplate(
  `
你是一个数学家，擅长使用中文回答数学相关的问题，当你不知道问题的答案时，你就回答不知道。
具体问题如下：
{input}
`
)
  .pipe(model)
  .pipe(parser); // 数学链

const otherChain = ChatPromptTemplate.fromTemplate(
  `
你是一个AI助手，你会使用中文回答以下问题。
具体问题如下：
{input}
`
)
  .pipe(model)
  .pipe(parser); // 其它链

// 4. 分类链：输出“数学”/“物理”/“其它”
const classifyChain = ChatPromptTemplate.fromMessages([
  [
    "system",
    `
你是一个仅输出标签的分类器。
任务：将用户问题分类为以下三类之一：数学、物理、其它。
规则：
- 只能输出且必须输出其中一个：数学 或 物理 或 其它
- 不要输出任何解释、理由、引号、标点、前后缀或额外空行
- 若不确定，输出：其它
    `.trim(),
  ],
  [
    "human",
    `
问题：
{input}

现在输出分类（仅一个词）：`.trim(),
  ],
])
  .pipe(model)
  .pipe(parser); // 分类链

// 5. TODO
// 分派器：根据 topic 路由到对应子链
const answerBranch = RunnableBranch.from([
  [
    (input) => input.topic.includes("物理"),
    RunnableLambda.from((x) => {
      console.log("执行了物理链");
      return x;
    }).pipe(physicsChain),
  ],
  [
    (input) => input.topic.includes("数学"),
    RunnableLambda.from((x) => {
      console.log("执行了数学链");
      return x;
    }).pipe(mathChain),
  ],
  otherChain,
]);

// 6. 最终组合 - 最终给外部使用的，就是这个 finalChain
const finalChain = RunnableSequence.from([
  // 当前的入参数：{ input: "什么是经典力学？" } --> { input: "什么是经典力学？" : topic: "物理"}
  RunnablePassthrough.assign({
    topic: classifyChain,
  }),
  answerBranch,
]);

// 测试
const run = async () => {
  console.log(await finalChain.invoke({ input: "什么是经典力学？" }));
  console.log(
    await finalChain.invoke({ input: "对 y = x 求导的结果是多少？" })
  );
  console.log(await finalChain.invoke({ input: "你好，你是谁？" }));
};
run();
